import torch
import torchvision
from torch import nn

# 或直接使用 'DEFAULT'（自动选择最新权重）
vgg16 = torchvision.models.vgg16(weights='DEFAULT')

# 保存方式1：模型结构+模型参数
torch.save(vgg16, 'vgg16_method1.pth')
# 保存方式2：模型参数（官方推荐）
torch.save(vgg16.state_dict(), 'vgg16_method2.pth')


# 陷进
# 陷阱1：保存的模型，如果想在代码中加载，必须使用与保存时相同的类结构
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)

    def forward(self, x):
        x = self.conv1(x)
        return x


net = Net()
# 方式一保存
torch.save(net, 'net_method1.pth')